#include <sstream>
#include <iomanip>

#include "base_util/utils.h"
#include "facefeature_predictor.h"
#include <opencv2/dnn.hpp>

namespace ai {


static void dump_tensor_attr(rknn_tensor_attr *attr) {
  printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
          "zp=%d, scale=%f\n",
          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}
static inline float fast_exp(float x) {
  // return exp(x);
  union
  {
    uint32_t i;
    float f;
  } v;
  v.i = (12102203.1616540672 * x + 1064807160.56887296);
  return v.f;
}
static float DeQnt2F32(int8_t qnt, int zp, float scale) {
  return ((float)qnt - (float)zp) * scale;
}

RKFacefeaturePredictor::RKFacefeaturePredictor(RKPackPredictor* model, RKModelManager* manager, LogInfo *lg):
  RKPredictor(model, manager, lg) { }


int RKFacefeaturePredictor::start_postprocess_thread_imp() {
  int ret = 0;
  std::pair<std::pair<int,int>,int> cur_block;
  // stop_flag2 = common::CameraGrabing;
  // while (stop_flag2 == common::CameraGrabing) {

  //   run_mtx2.lock();
  //   if (post_data2.empty()) {run_mtx2.unlock(); std::this_thread::sleep_for(std::chrono::milliseconds(3));continue;}
  //   cur_block = post_data2.front();
  //   post_data2.pop();
  //   run_mtx2.unlock();

  //   int loop_idx = cur_block.first.second;
  //   std::vector<BaseInfo*> det_infos;
  //   ret = det_hb_->PostProcess(det_infos, cur_block.first.first, this, cur_block.second);  //处理图片
  //   run_mtx3.lock();
  //   out_data3[loop_idx].push(det_infos);  //将处理后的数据返回到det_infos,加入队列中
  //   run_mtx3.unlock();
  // }
  // stop_flag2 = common::CameraOpened;
  return 0;
}

int RKFacefeaturePredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret;
  // printf("[liufeng debug] RKFacefeaturePredictor::RunDet:\n");
  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  /*
  if (!mdl_rk->cfg->transforms->run(blob->img, img_blob, mdl_rk->cfg)) { 
    printf("transforms->run fail \n");
    return model_image_channels_check_error;
  }
  */
  // std::cout << "img_blob.img w:" << img_blob.img.cols << " h:" << img_blob.img.rows << std::endl;
  // cv::imwrite("img_blob.img.jpg",img_blob.img);
  // cv::imwrite("blobimg.jpg",blob->img);

  img_blob.ori_im_shape = blob->ori_im_shape;

  rknn_tensor_attr output_attrs[mdl_rk->output_count];
  memset(output_attrs, 0, sizeof(output_attrs));
  for (int i = 0; i < mdl_rk->output_count; i++) {
    output_attrs[i].index = i;
    ret = rknn_query(mdl_rk->ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
    // dump_tensor_attr(&(output_attrs[i]));
  } 

  // mdl_rk->inputs[0].buf = (void *)(img_blob.img.data);
  mdl_rk->inputs[0].buf = (void *)(blob->img.data);
  rknn_inputs_set(mdl_rk->ctx, mdl_rk->input_count, &mdl_rk->inputs[0]);

  rknn_core_mask core_mask = RKNN_NPU_CORE_0;
  if(mdl_rk->cfg->gpu_id==1) { core_mask = RKNN_NPU_CORE_1; }
  if(mdl_rk->cfg->gpu_id==2) { core_mask = RKNN_NPU_CORE_2; }
  ret = rknn_set_core_mask(mdl_rk->ctx, core_mask);
  ret = rknn_run(mdl_rk->ctx, NULL);
  if (ret != 0) {printf("rknn_run fail %d", ret); return -1;}
  ret = rknn_outputs_get(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0], NULL);
  if (ret != 0) {printf("rknn_outputs_get fail %d", ret); return -1;}
 
  // 后处理部分
  std::vector<float> out_scales;
  std::vector<int32_t> out_zps;
  for (int i = 0; i < mdl_rk->output_count; ++i) {
    out_scales.push_back(output_attrs[i].scale);
    out_zps.push_back(output_attrs[i].zp);
  }

  float* output_tensors[mdl_rk->output_count];
  for (int j = 0; j < mdl_rk->output_count; j++)
  {

      int8_t *pblob = (int8_t *)mdl_rk->outputs[j].buf;
      float* pf=(float*)malloc(output_attrs[j].n_elems*sizeof(float));
      
      for(int i=0;i<output_attrs[j].n_elems;i++)
      {
          pf[i]=DeQnt2F32(pblob[i], out_zps[j], out_scales[j]);               
      }
      output_tensors[j]=pf;

  }

  DetInfo* di = new DetInfo();
  for(int i=0;i<output_attrs[0].n_elems;i++)
  {
    di->kps.push_back(output_tensors[0][i]);
  }
  det_infos.emplace_back(di);

  for(int i=0;i<mdl_rk->output_count;i++)
  {
      free(output_tensors[i]);
  }
  rknn_outputs_release(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0]);

  return 0;

}

int RKFacefeaturePredictor::PostProcess(std::vector<BaseInfo*>& det_infos, 
                int cur_block,
                RKModelManager* mng, 
                int md_idx) {

  return 0;
}

void RKFacefeaturePredictor::ParseTensor(float* data, int layer, 
  std::vector<cv::Rect>& boxs, std::vector<float>& scores, std::vector<int>& labels,
  std::vector<int> ori_shape, std::vector<int> feteature_shape, ModelConfig* cfg) {
    int num_classes = cfg->label_list.size();
    int stride = cfg->strides[layer];
    int num_pred = num_classes + 4 + 1;
    if (cfg->anchors_table.empty()) {
      spdlog::get("logger")->info("ERROR. RKFacefeaturePredictor::ParseTensor anchors_table is empty.");
      return;
    }
    std::vector<std::vector<int>> &anchors = cfg->anchors_table[layer];
    float w_scale = ori_shape[1] * 1.0 / cfg->input_shape[1];
    float h_scale = ori_shape[0] * 1.0 / cfg->input_shape[0];
    // spdlog::get("logger")->info("num_classes: {}, stride: {}, num_pred: {}, w_scale: {}, h_scale: {}", num_classes, stride, num_pred, w_scale, h_scale);

    int anchor_num = anchors.size();
    for (int h = 0; h < feteature_shape[0]; h++) {
      for (int w = 0; w < feteature_shape[1]; w++) {
        for (int k = 0; k < anchor_num; k++) {
          float *cur_data = data + k * num_pred;
          // if (h == 0 && w == 0 ) {printf("%f, %f, %f, %f, %f, %f, %f, %f\n", cur_data[0],cur_data[1],cur_data[2],cur_data[3],cur_data[4],cur_data[5],cur_data[6],cur_data[7]);}
          
          int id = std::distance(cur_data + 5, std::max_element(cur_data + 5, cur_data + 5 + num_classes));
          double x1 = 1 / (1 + std::exp(-cur_data[4])) * 1;
          double x2 = 1 / (1 + std::exp(-cur_data[id + 5]));
          double confidence = x1 * x2;
          // if (1 ) {printf("id:%d, 4:%f, id+5:%f, x1:%f, x2:%f, confidence:%f\n", id, cur_data[4], cur_data[id + 5], x1, x2, confidence);}
          if (confidence < cfg->draw_threshold) {continue;}
          // spdlog::get("logger")->info("confidence:{}", confidence);

          double cx = ((1.0 / (1.0 + std::exp(-cur_data[0]))) * 2 - 0.5 + w) * stride;
          double cy = ((1.0 / (1.0 + std::exp(-cur_data[1]))) * 2 - 0.5 + h) * stride;
          double ww = std::pow((1.0 / (1.0 + std::exp(-cur_data[2]))) * 2, 2) * anchors[k][0];
          double hh = std::pow((1.0 / (1.0 + std::exp(-cur_data[3]))) * 2, 2) * anchors[k][1];

          double xmin = (cx - ww / 2.0)* w_scale;
          double ymin = (cy - hh / 2.0)* h_scale;
          double xmax = (cx + ww / 2.0)* w_scale;
          double ymax = (cy + hh / 2.0)* h_scale;

          if (xmax <= 0 || ymax <= 0) { continue; }
          if (xmin > xmax || ymin > ymax) { continue; }
          if (xmin < 0) { xmin = 0;}
          if (ymin < 0) { ymin = 0;}
          if (xmax >= ori_shape[1]) { xmax = ori_shape[1]-1;}
          if (ymax >= ori_shape[0]) { ymax = ori_shape[0]-1;}

          boxs.emplace_back(cv::Rect(xmin, ymin, xmax-xmin, ymax-ymin));
          scores.emplace_back((float)confidence);
          labels.emplace_back(id);
        }
        data = data + num_pred * anchors.size();
      }
    }
}


// int RKFacefeaturePredictor::get_tensor_hw(hbDNNTensor* tensor, int *height, int *width) {
//   int h_index = 0;
//   int w_index = 0;
//   if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NHWC) {
//     h_index = 1;
//     w_index = 2;
//   }
//   else if (tensor->properties.tensorLayout == HB_DNN_LAYOUT_NCHW) {
//     h_index = 2;
//     w_index = 3;
//   }
//   else { return -1; }
//   *height = tensor->properties.validShape.dimensionSize[h_index];
//   *width = tensor->properties.validShape.dimensionSize[w_index];
//   return 0;
// }


}  // namespace hb
